import numpy as np
import taichi as ti
from scipy import signal as Signal
import tensorly as tensorly
from tensorly.decomposition import symmetric_parafac_power_iteration
from tqdm import tqdm
from Utils import debug_output, numpy_to_taichi_field
import os

class Kernel():
    def __init__(
        self,
        kernal_path = './KernelDataBase',
        max_rank = 4,
        jacobi_iters = 128,
        kappa_temperature = 10.1,
        kappa_density = 10.1,
        kappa_viscosity = 10.1,
        delta_time = 0.03,
        delta_scale = 4,
        debug_mode = False

    ) -> None:
        self.kernel_path = kernal_path
        self.max_rank = max_rank
        self.jacobi_iters = jacobi_iters
        self.kappa_temperature = kappa_temperature
        self.kappa_density = kappa_density
        self.kappa_viscosity = kappa_viscosity
        self.delta_time = delta_time
        self.delta_scale = delta_scale

        self.velocity_initial_poisson_weight_2D = None
        self.velocity_initial_poisson_vector_2D = None
        self.velocity_poisson_weight_2D = None
        self.velocity_poisson_vector_2D = None
        self.velocity_diffuse_weight_2D = None
        self.velocity_diffuse_vector_2D = None
        self.density_diffuse_weight_2D = None
        self.density_diffuse_vector_2D = None
        self.temperature_diffuse_weight_2D = None
        self.temperature_diffuse_vector_2D = None
        self.debug_mode = debug_mode

        # print('before main : ',self.density_diffuse_weight_2D)
        if not self.load_kernel():
            print("No Initial Kernel --- Start Generate", end='\n')
            self.generate_kernel()
            self.output_kernel()

        # print('In main : ',self.density_diffuse_weight_2D)

        self.to_taichi()
        pass


    def to_taichi(self):
        self.velocity_initial_poisson_weight_2D = numpy_to_taichi_field(self.velocity_initial_poisson_weight_2D)
        self.velocity_initial_poisson_vector_2D = numpy_to_taichi_field(self.velocity_initial_poisson_vector_2D)
        self.velocity_poisson_weight_2D = numpy_to_taichi_field(self.velocity_poisson_weight_2D)
        self.velocity_poisson_vector_2D = numpy_to_taichi_field(self.velocity_poisson_vector_2D)
        self.velocity_diffuse_weight_2D = numpy_to_taichi_field(self.velocity_diffuse_weight_2D)
        self.velocity_diffuse_vector_2D = numpy_to_taichi_field(self.velocity_diffuse_vector_2D)
        self.density_diffuse_weight_2D = numpy_to_taichi_field(self.density_diffuse_weight_2D)
        self.density_diffuse_vector_2D = numpy_to_taichi_field(self.density_diffuse_vector_2D)
        self.temperature_diffuse_weight_2D = numpy_to_taichi_field(self.temperature_diffuse_weight_2D)
        self.temperature_diffuse_vector_2D = numpy_to_taichi_field(self.temperature_diffuse_vector_2D)
        # print(self.density_diffuse_weight_2D)
    def GenPoissonKernelInitial2D(self, Beta):
        KernelMid = self.jacobi_iters - 1
        KernelSize = self.jacobi_iters * 2 - 1

        PoissonKernel = np.zeros(
            shape = [KernelSize, KernelSize],
            dtype = np.double
        )
        PoissonKernel[KernelMid, KernelMid] = 1
        ConvKernel = np.array(
            [
                [0.      , 1 / Beta, 0       ],
                [1 / Beta, 0       , 1 / Beta],
                [0.      , 1 / Beta, 0       ]
            ], dtype=np.double
        )

        for _ in range(self.jacobi_iters):
            PoissonKernelTemp = Signal.convolve2d(
                PoissonKernel, ConvKernel, mode='same'
            )
            PoissonKernel = PoissonKernelTemp

        return tensorly.tensor(PoissonKernel)

    def GenPoissonKernel2D(self, Alpha, Beta):
        KernelMid = self.jacobi_iters - 1
        KernelSize = self.jacobi_iters * 2 - 1

        PoissonKernel = np.zeros(
            shape = [KernelSize, KernelSize],
            dtype=np.double
        )
        ConvKernel = np.array(
            [
                [0.      , 1 / Beta, 0       ],
                [1 / Beta, 0       , 1 / Beta],
                [0.      , 1 / Beta, 0       ]
            ], dtype=np.double
        )
        for _ in range(self.jacobi_iters):
            PoissonKernelTemp = Signal.convolve2d(
                PoissonKernel, ConvKernel, mode='same'
            )
            PoissonKernel = PoissonKernelTemp
            PoissonKernel[KernelMid, KernelMid] += Alpha / Beta

        return tensorly.tensor(PoissonKernel)

    def GenPoissonKernelInverseInitial2D(self):
        Alpha = -self.delta_scale * self.delta_scale
        Beta = 4
        return self.GenPoissonKernelInitial2D(Beta)

    def GenPoissonKernelInverse2D(self):
        Alpha = -self.delta_scale * self.delta_scale
        Beta = 4
        return self.GenPoissonKernel2D(Alpha, Beta)

    def GenPoissonKernelForward2D(self, Kappa):
        Alpha = self.delta_scale * self.delta_scale / (Kappa * self.delta_time)
        Beta = 4 + Alpha
        return self.GenPoissonKernel2D(Alpha, Beta)

    def output_kernel(self):
        file_name = self.get_kernel_name()
        file_path = os.path.join(self.kernel_path, file_name)

        if os.path.exists(self.kernel_path) == False :
            os.mkdir(self.kernel_path)

        np.savez(
            file_path,
            velocity_initial_poisson_weight_2D = self.velocity_initial_poisson_weight_2D,
            velocity_initial_poisson_vector_2D = self.velocity_initial_poisson_vector_2D,
            velocity_poisson_weight_2D = self.velocity_poisson_weight_2D,
            velocity_poisson_vector_2D = self.velocity_poisson_vector_2D,
            velocity_diffuse_weight_2D = self.velocity_diffuse_weight_2D,
            velocity_diffuse_vector_2D = self.velocity_diffuse_vector_2D,
            density_diffuse_weight_2D = self.density_diffuse_weight_2D,
            density_diffuse_vector_2D = self.density_diffuse_vector_2D,
            temperature_diffuse_weight_2D = self.temperature_diffuse_weight_2D,
            temperature_diffuse_vector_2D = self.temperature_diffuse_vector_2D,
        )
        pass

    def get_kernel_name(self):
        return '{}_{}_{}_{}_{}_{}_{}_{}'.format(
            self.max_rank,
            self.jacobi_iters,
            self.kappa_temperature,
            self.kappa_density,
            self.kappa_viscosity,
            self.kappa_viscosity,
            self.delta_time,
            self.delta_scale
        )
        pass

    def load_kernel(self):

        file_name = self.get_kernel_name()
        file_path = os.path.join(self.kernel_path, file_name+'.npz')
        # print('Kernel Path : {}'.format(file_path))
        if os.path.exists(file_path):
            npzfile = np.load(file_path)
            # print(npzfile)
            # print(npzfile.files)
            self.velocity_initial_poisson_weight_2D = npzfile['velocity_initial_poisson_weight_2D']
            self.velocity_initial_poisson_vector_2D = npzfile['velocity_initial_poisson_vector_2D']
            self.velocity_poisson_weight_2D = npzfile['velocity_poisson_weight_2D']
            self.velocity_poisson_vector_2D = npzfile['velocity_poisson_vector_2D']
            self.velocity_diffuse_weight_2D = npzfile['velocity_diffuse_weight_2D']
            self.velocity_diffuse_vector_2D = npzfile['velocity_diffuse_vector_2D']
            self.density_diffuse_weight_2D = npzfile['density_diffuse_weight_2D']
            self.density_diffuse_vector_2D = npzfile['density_diffuse_vector_2D']
            self.temperature_diffuse_weight_2D = npzfile['temperature_diffuse_weight_2D']
            self.temperature_diffuse_vector_2D = npzfile['temperature_diffuse_vector_2D']
            # print('Loaded : '.format(npzfile['density_diffuse_weight_2D']))
            # print('Loaded : '.format(self.density_diffuse_weight_2D))
            debug_output(self.temperature_diffuse_weight_2D, debug_mode=self.debug_mode)
            debug_output(self.temperature_diffuse_vector_2D, debug_mode=self.debug_mode)
            debug_output(self.temperature_diffuse_weight_2D.shape, debug_mode=self.debug_mode)
            debug_output(self.temperature_diffuse_vector_2D.shape, debug_mode=self.debug_mode)


            return True
        else:
            return False

    def generate_kernel(self):

        with tqdm(total=5, ncols=60, desc='Kernel Generating') as tbar:
            PoissonKernel = self.GenPoissonKernelInverseInitial2D()
            PoissonWeights, PoissonVectors = symmetric_parafac_power_iteration(PoissonKernel, rank = self.max_rank, n_iteration=10, n_repeat=10)
            PoissonVectors = tensorly.transpose(PoissonVectors)
            self.velocity_initial_poisson_weight_2D = PoissonWeights
            self.velocity_initial_poisson_vector_2D = PoissonVectors
            tbar.update(1)

            PoissonKernel = self.GenPoissonKernelInverse2D()
            PoissonWeights, PoissonVectors = symmetric_parafac_power_iteration(PoissonKernel, rank = self.max_rank, n_iteration=10, n_repeat=10)
            PoissonVectors = tensorly.transpose(PoissonVectors)
            self.velocity_poisson_weight_2D = PoissonWeights
            self.velocity_poisson_vector_2D = PoissonVectors
            tbar.update(1)

            PoissonKernel = self.GenPoissonKernelForward2D(self.kappa_viscosity)
            PoissonWeights, PoissonVectors = symmetric_parafac_power_iteration(PoissonKernel, rank = self.max_rank)
            PoissonVectors = tensorly.transpose(PoissonVectors)
            self.velocity_diffuse_weight_2D = PoissonWeights
            self.velocity_diffuse_vector_2D = PoissonVectors
            tbar.update(1)

            PoissonKernel = self.GenPoissonKernelForward2D(self.kappa_density)
            PoissonWeights, PoissonVectors = symmetric_parafac_power_iteration(PoissonKernel, rank = self.max_rank)
            PoissonVectors = tensorly.transpose(PoissonVectors)
            self.density_diffuse_weight_2D = PoissonWeights
            self.density_diffuse_vector_2D = PoissonVectors
            tbar.update(1)

            PoissonKernel = self.GenPoissonKernelForward2D(self.kappa_temperature)
            PoissonWeights, PoissonVectors = symmetric_parafac_power_iteration(PoissonKernel, rank = self.max_rank)
            PoissonVectors = tensorly.transpose(PoissonVectors)
            self.temperature_diffuse_weight_2D = PoissonWeights
            self.temperature_diffuse_vector_2D = PoissonVectors
            tbar.update(1)
            
            debug_output(self.temperature_diffuse_weight_2D, debug_mode=self.debug_mode)
            debug_output(self.temperature_diffuse_vector_2D, debug_mode=self.debug_mode)

        pass


if __name__ == '__main__':
    print("Kernel Testing")
    ti.init(arch=ti.gpu)
    kernel = Kernel(debug_mode = True)